import numpy as np

class SatelliteProcessor:
    def __init__(self, sat):
        self.sat = sat  # 위성 종류를 클래스 상태로 저장

    def quality_control(self, df):
        if self.sat in ['TROPOMI', 'GEMS_IUP']:
            df[df.qf < 0.75] = np.nan
            df[df.cf > 0.33] = np.nan
            df[df[f'{self.sat}_TCD'] >= 9.9e36] = np.nan
        elif self.sat in ['GEMS', 'GEMS_V3']:
            df[df.qf != 0] = np.nan
            df[df.cf > 0.33] = np.nan
            df[df[f'{self.sat}_TCD'] <= -1e29] = np.nan
        return df

    def select_prod(self, df, prod):
        if prod == 'TCD':
            col_varname = prod
        else:
            col_varname = 'VCD'

        if self.sat == 'TROPOMI':
            col_list = [f'{self.sat}_{col_varname}', f'{col_varname}_uncertainty', 'RS_TCD', 'RS_VCD', 'Pixel_area', 'Season', 'SZA']
        else:
            col_list = [f'{self.sat}_{col_varname}', f'RS_{col_varname}', 'Season', 'SZA']

        df_ext = df[col_list].dropna(axis=0)
        df_ext['Diff'] = df_ext[f'{self.sat}_{col_varname}'] - df_ext[f'RS_{col_varname}']
        df_ext['Rel_diff'] = (df_ext[f'{self.sat}_{col_varname}'] - df_ext[f'RS_{col_varname}']) / df_ext[f'RS_{col_varname}']
        df_ext['Hour'] = df_ext.index.round(freq='h').hour
        df_ext['dow'] = df_ext.index.day_of_week
        return df_ext

    @staticmethod
    def get_season(month):
        if 3 <= month <= 5:
            return 'MAM'
        elif 6 <= month <= 8:
            return 'JJA'
        elif 9 <= month <= 11:
            return 'SON'
        else:
            return 'DJF'